import re
import socketserver as SocketServer
import time, sys
import argparse
import NLAgent
import select
from select import poll
import threading
from socket import *      
from collections import deque
import resource

parser = argparse.ArgumentParser(parents=[NLAgent.parser], description='Transport Layer Agent: Server')
server_args = parser.add_argument_group('Server Behavior')
resource.setrlimit(resource.RLIMIT_NOFILE, (65536, 65536))
(soft, hard) = resource.getrlimit(resource.RLIMIT_NOFILE)
print("Limits are ", soft, " and ", hard)

poll = select.poll()
DBfilename=""

def cleanup():
  print("Done")

STIME = 1
CTIME = 0.05
ERTHRESH = 0
MAXJOBS = 100000
MAXSERVERS = 10
MAXCLIENTS = 10
MAXACCRCV = 2
MAXWQUEUES = 10

wait_cnt = 0
wait_queue = []
wait_queue_lock = []
start_time = time.time()
connections = dict()
connids = dict()
errors = dict()
queue_servers = deque([])
queue_clients = deque([])
queue_acceptsendreceive = deque([])
select_inputs = []
listeners = []
sock2listener = dict()
sock2conn = dict()
connreceived = dict()
fd2socket = dict()
server_lock = threading.Lock()
client_lock = threading.Lock()
socket_lock = threading.Lock()
conn_lock = threading.Lock()
acceptsendreceive_lock = threading.Lock()
wait_lock = threading.Lock()

class Job():
    def __init__(self, cid, jtype, value, socket, host, port, lsocket):
      self.cid = cid
      self.type = jtype
      self.value = value
      self.socket = socket
      self.host = host
      self.port = port
      self.lsocket = lsocket
    

class ServData():
    def __init__(self, cid, host, port, time):
      self.cid = cid
      self.host = host
      self.port = port
      self.start = time
      self.conns = dict()

class ConnID():
    def __init__(self, cid, server, sport, client, cport, time):
      self.cid = cid
      self.server = server
      self.sport = sport
      self.client = client
      self.cport = cport
      self.start = time
      self.found = 0
      self.state = -1

class Event():
    def __init__(self, type, value):
        self.type = type
        self.value = value

class ConnData():
    def __init__(self):
        self.events = deque([])


def handle_next_event(c, socket, caller):
    global queue_acceptsendreceive, connections, select_inputs, acceptsendreceive_lock, wait_queue, wait_cnt, wait_lock, wait_queue_lock, start_time, connreceived, fd2socket

    print("HNE for ", c, " event len ", len(connections[c].events), " called from ", caller)
    # Check if there are no more events
    if len(connections[c].events) == 0:
      socket_lock.acquire()
      if socket in select_inputs:
        select_inputs.remove(socket)
        poll.unregister(socket.fileno())
        print("SI removed and will close ", socket)
        connids[c].state = 1
        socket.close()              
      socket_lock.release()
      return
    e = connections[c].events[0]
    print("HNE for ", c, " event type ", e.type)
    if e.type == 'RECV':
      pass
      # RECV events will be removed only when they are handled
    elif e.type == 'WAIT':
      # WAIT events are removed, queued and sorted
      # We use multiple queues to keep the queue size small
      # We only queue up an event if its duration is longer than
      # 10 ms
      try:
        connections[c].events.popleft()
      except Exception as inst:
        print("Exception in HNE for ",c," is ", type(inst)) 
      if e.value > 10:
        now = time.time() - start_time
        job = Job(c,'WAIT', now + e.value/1000.0, socket, 0, 0, -1)
        wait_lock.acquire()
        wait_cnt += 1
        if wait_cnt == MAXWQUEUES:
          wait_cnt = 0
        cnt = wait_cnt
        wait_lock.release()
        wait_queue_lock[cnt].acquire()
        wait_queue[cnt].append(job)
        wait_queue[cnt].sort(key=lambda job: job.value)
        wait_queue_lock[cnt].release()
        print("HNEW: appended WAIT event for ", c, " to ", cnt, " value ", job.value, " event value ", e.value/1000.0, " time now ", now);
      else:
        handle_next_event(c, socket, "handle")
    else:
      # SEND events are removed and queued
      connections[c].events.popleft()
      job = Job(c,'SEND',e.value, socket, 0, 0, -1)
      acceptsendreceive_lock.acquire()
      queue_acceptsendreceive.append(job)
      acceptsendreceive_lock.release()
      print("HNEW: appended SEND event for ", c);

class WaitThread(threading.Thread):


  def run (self):
    global wait_queue, wait_cnt, wait_lock, wait_queue_lock, start_time
    while True:
      job = None
      for i in range(0, MAXWQUEUES):
        # Do a pre-check without acquiring the lock
        if len(wait_queue[i]):
          job = None
          wait_queue_lock[i].acquire()
          if len(wait_queue[i]):
            job = wait_queue[i][0]
            if(time.time() - start_time >= job.value):
              print("WAIT: next job is up ", job.value, " queue ", i)
              job = wait_queue[i].pop(0)
            else:
              job = None
          wait_queue_lock[i].release()
          if job != None:
           handle_next_event(job.cid, job.socket, "wait")




class AcceptSendReceiveThread(threading.Thread):


  def run (self):
    global listeners, queue_acceptsendreceive, sock2conn, connections, select_inputs, acceptsendreceive_lock, start_time, fd2socket, poll
    while True:
      job = None
      acceptsendreceive_lock.acquire()
      try:
        job = queue_acceptsendreceive.popleft()
      except:
        pass
      acceptsendreceive_lock.release()
      if job != None:
        print("Queue ASR popped job ", job.type, " for conn ", job.cid)
        if job.type == 'ACCEPT':
          print("AR: Will find connection ", job.host, " on port ", job.port)
          s = sock2listener[job.lsocket]
          for c in s.conns:
            if s.conns[c].client == job.host and s.conns[c].cport == job.port and s.conns[c].found == 0:
              print("Got connection ", c, " client ", s.conns[c].client, " port ", s.conns[c].cport, " on socket ", job.socket, " server IP ", connids[c].server, " sport ", connids[c].sport)
              connids[c].state = 0
              sock2conn[job.socket] = c
              sock2listener[job.lsocket].conns[c].found = 1
              socket_lock.acquire()
              select_inputs.append(job.socket)
              connreceived[c] = 0
              print("SI: appended ", job.socket)
              socket_lock.release()
              handle_next_event(c, job.socket, " AR thread acc")
        elif job.type == 'SEND':
          print("AR: Handling SEND for ", job.cid)
          try:
            job.socket.sendall(b'a'*job.value) 
            print("Sent data for ", job.cid, " sent bytes ", job.value)
          except Exception as inst:
            print("Exception ", self.getName(), " send failed ", type(inst))
          handle_next_event(job.cid, job.socket, " AR thread send")
        else:
          pass

class StartServersThread (threading.Thread):

   def run (self):
     global queue_servers, server_lock, sock2listener, socket_lock, conn_lock, start_time, fd2socket, poll

     while (len(queue_servers)):
       s = None
       server_lock.acquire()
       if (len(queue_servers) and time.time() - start_time >= queue_servers[0].start - STIME):
         s = queue_servers.popleft()         
         print(self.getName(), " popped for ", s.cid, " start at ", s.start, " now ", time.time() - start_time,  " starting server for ", s.host, " and ", s.port, " time ", time.time(), " now ", time.time() - start_time, " start time ", start_time, " left ", len(queue_servers))
       server_lock.release()
       if (s != None):
         serv = socket(AF_INET,SOCK_STREAM)    
         serv.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
         try: 
           serv.bind((s.host, s.port))  
         except:
           print("Exception, could not bind ", s.host, " ", s.port)
           return
         serv.listen(5)  
         print("Server ", s.host, " and ", s.port, " started at ", time.time() - start_time)
         socket_lock.acquire()
         select_inputs.append(serv)
         poll.register(serv.fileno(), select.POLLIN)
         fd2socket[serv.fileno()] = serv
         print("SI: appended ", serv)
         listeners.append(serv)
         socket_lock.release()
         conn_lock.acquire()
         sock2listener[serv] = s
         conn_lock.release()


class StartClientsThread (threading.Thread):

   def run (self):
     global queue_clients, client_lock, socket_lock, connections, start_time, fd2socket, poll

     while True:
       s = None
       #print("Client thread ", self.getName(), " time ", time.time() - start_time)
       client_lock.acquire()
       if (len(queue_clients) and time.time() - start_time >= queue_clients[0].start - CTIME):
         s = queue_clients.popleft()         
         print(self.getName(), " popped for ", s.cid, " start at ", s.start, " now ", time.time() - start_time,  " starting client for ", s.client, " and ", s.cport)
       client_lock.release()
       #print("Client thread ", self.getName(), " done w CLIENT LOCK time ", time.time() - start_time)
       if (s != None):
         cl = socket(AF_INET,SOCK_STREAM)    
         cl.setblocking(0)
         cl.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
         try:
           cl.bind((s.client, s.cport))  
           print("Bound for ", s.cid)
         except:
           print("Exception, could not bind to port for ", s.cid)
           #print("Client thread ", self.getName(), " DONE w exception time ", time.time() - start_time)
           return
         #print("Client thread ", self.getName(), " ask for CLIENT LOCK time ", time.time() - start_time)
         client_lock.acquire()
         poll.register(cl.fileno(), select.POLLIN)
         poll.register(cl.fileno(), select.POLLOUT)
         client_lock.release()
         #print("Client thread ", self.getName(), " done with CLIENT LOCK time ", time.time() - start_time)
         fd2socket[cl.fileno()] = cl
         sock2conn[cl] = s.cid
         try:
           print("Trying to connect for ", s.cid, " socket ", cl)
           cl.connect((s.server, s.sport))
         except Exception as e:
           pass
         #print("Client thread ", self.getName(), " DONE time ", time.time() - start_time)


         
def init_servers(db_filename, master_start_time, ip_list, stime):
  servers = {}

  i = 0  
  all_servers = dict()
  global listeners, connections
  
  # Remember all connections we have to set up
  try:
    db = NLAgent.db_connect(db_filename)
    start_time = 0
    cac = 0
    cas = 0
    for connection in db.get_list_of_connections():
      if (start_time == 0):
        start_time = connection['start_time']
      dst = connection['dst'].strip()
      src = connection['src'].strip()
      cid = connection['conn_id']
      src_invalid = re.search('(\.0|\.255)$', src)
      dst_invalid = re.search('(\.0|\.255)$', dst)
      if src_invalid or dst_invalid or connection['dst_port'] == 22:
        print("Skipping ", cid, " src ", src, " and ", dst, " port ", connection['dst_port'])
        continue
      print("Creating ", cid)
      connections[cid] = ConnData()
      HOST,PORT = connection['dst'].strip(), connection['dst_port']
      c = ConnID(cid, HOST, PORT, connection['src'], connection['src_port'], connection['start_time'] - start_time)
      connids[cid] = c
      if connection['dst'].strip() in ip_list:
        cas += 1
        sid = str(HOST) + "-" + str(PORT)
        if (sid not in all_servers):
          all_servers[sid] = i
          s = ServData(sid, HOST, PORT, connection['start_time'] - start_time)
          s.conns[cid] = c
          queue_servers.append(s)
          print("Will start server ", sid, " at ", connection['start_time'] - start_time)
          i += 1
        else:
          queue_servers[all_servers[sid]].conns[cid] = c
      else: 
        cac += 1
        queue_clients.append(c)
        errors[cid] = 0
    print("Will open ", cac, " connections and ", cas, " server connections")
          
  except Exception as e:
    print("Exception ", e)
    return

  # Remember all events
  for event in db.get_all_events():
    cid = event['conn_id']
    if cid in connections:
      if event['actor_ip'] in ip_list:
        connections[cid].events.append(Event(event['state'], event['value']))
      elif event['state'] == 'SEND':
        connections[cid].events.append(Event('RECV', event['value']))
    else:
      print("No connection ", cid)

def run(db_filename, master_start_time, ip_list, stime):
  global select_inputs,listeners, connections, queue_acceptsendreceive, acceptsendreceive_lock, sock2listener, sock2conn, start_time, wait_queue, wait_queue_lock, fd2socket, poll, queue_clients, client_lock, errors

  start_time = time.time()
  select_inputs = []
  select_outputs = []
  # Start threads that will open connections
  print("Will now start servers")
  startServers = []
  startClients = []
  acceptSendReceivers = []
  waiters = []
  for i in range(0, MAXWQUEUES):
    wait_queue.append([])
    l = threading.Lock()
    wait_queue_lock.append(l)
  for i in range(0,MAXSERVERS):
    startServers.append(StartServersThread())
    startServers[i].start()
  for i in range(0,MAXCLIENTS):
    startClients.append(StartClientsThread())
    startClients[i].start()
  for i in range(0,MAXACCRCV):
    acceptSendReceivers.append(AcceptSendReceiveThread())
    acceptSendReceivers[i].start()
  for i in range(0, MAXWQUEUES):
    waiters.append(WaitThread())
    waiters[i].start()

  while True:
    loop_time = time.time()
    try: 
      socket_lock.acquire()
      socketevents = poll.poll()
      socket_lock.release()
      for (fd,ev) in socketevents:
        socket = fd2socket[fd]
        print("Socket ", socket, " event ", ev)
        if ev != select.POLLIN and ev != select.POLLOUT:
          print("Some exception happened, see if we should try again ", socket)
          cid = sock2conn[socket]
          if errors[cid] > ERTHRESH:
            print("SI removed and will close ", socket)
            poll.unregister(fd)
            connids[sock2conn[socket]].state = 2
            socket.close()
          else:
            errors[cid] += 1
            client_lock.acquire()
            queue_clients.appendleft(connids[cid])          
            client_lock.release()
        else:
          if ev == select.POLLIN:
            if socket in listeners:
              (conn, (host, port)) = socket.accept()
              poll.register(conn.fileno(), select.POLLIN)
              fd2socket[conn.fileno()] = conn
              print("AR: Will find connection ", host, " on port ", port)
              s = sock2listener[socket]
              for c in s.conns:
                if s.conns[c].client == host and s.conns[c].cport == port and s.conns[c].found == 0:
                  print("Got connection ", c, " client ", s.conns[c].client, " port ", s.conns[c].cport, " on socket ", conn, " server IP ", connids[c].server, " sport ", connids[c].sport)
                  connids[c].state = 0
                  sock2conn[conn] = c
                  sock2listener[socket].conns[c].found = 1
                  socket_lock.acquire()
                  select_inputs.append(conn)
                  connreceived[c] = 0
                  print("SI: appended ", conn)
                  socket_lock.release()
                  handle_next_event(c, conn, " poll acc")
                  break
            else:
              data = ''
              try:
                data = socket.recv(8192)
                print(time.time(), " ", time.time() - start_time, "Received ", len(data))
                cid = sock2conn[socket]
              except:
                print("Exception: no socket ", socket)
            
              if len(data):
                print("Received ", len(data), " for ", cid, " events ", len(connections[cid].events))
                # Wait for the thread to fill in accept info
                unsynced = True
                while unsynced:
                  try:
                    connreceived[cid] += len(data)
                    unsynced = False
                  except:
                    pass
                print("Connreceived for ", cid, " is ", connreceived[cid])
                if len(connections[cid].events):
                  print("Expecting for ", cid, " to receive ", connections[cid].events[0].value)
                else:
                  print("No events left for ", cid)
                try: 
                  if (connreceived[cid] == connections[cid].events[0].value):
                    print("I've received everything for ", cid)
                    connreceived[cid] = 0
                    connections[cid].events.popleft()
                    handle_next_event(cid, socket, "receive")
                except:
                  print("Index out of range for ", cid, " events ", len(connections[cid].events))
              else:
                pass
          elif ev == select.POLLOUT:
           print("Socket ", socket, " event POLLOUT")
            # Some connection has succeeded
           socket_lock.acquire()
           select_inputs.append(socket)
           cid = sock2conn[socket]
           connreceived[cid] = 0
           print("SI appended ", socket)
           poll.modify(fd, select.POLLIN)
           socket_lock.release()
           print("Connection succeeded for ", cid)
           connids[cid].state = 0
           handle_next_event(cid, socket, " connect")


    except KeyboardInterrupt:
      print("Keyboard interrupt")
      for c in connids:
        print("Connection ", c, " state ", connids[c].state, " server ", connids[c].server, " port ", connids[c].sport)
      break
  
if __name__ == "__main__":
  args = parser.parse_args()  
  
  
  
    
